//----------------------------------------------------------------------------//
//                                                                            //
// ozz-animation is hosted at http://github.com/guillaumeblanc/ozz-animation  //
// and distributed under the MIT License (MIT).                               //
//                                                                            //
// Copyright (c) Guillaume Blanc                                              //
//                                                                            //
// Permission is hereby granted, free of charge, to any person obtaining a    //
// copy of this software and associated documentation files (the "Software"), //
// to deal in the Software without restriction, including without limitation  //
// the rights to use, copy, modify, merge, publish, distribute, sublicense,   //
// and/or sell copies of the Software, and to permit persons to whom the      //
// Software is furnished to do so, subject to the following conditions:       //
//                                                                            //
// The above copyright notice and this permission notice shall be included in //
// all copies or substantial portions of the Software.                        //
//                                                                            //
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR //
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,   //
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL    //
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER //
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING    //
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER        //
// DEALINGS IN THE SOFTWARE.                                                  //
//                                                                            //
//----------------------------------------------------------------------------//

#include "SkrAnim/ozz/ik_two_bone_job.h"

#include <cassert>

#include "SkrAnim/ozz/base/log.h"
#include "SkrAnim/ozz/base/maths/math_ex.h"
#include "SkrAnim/ozz/base/maths/simd_quaternion.h"

using namespace ozz::math;

namespace ozz
{
namespace animation
{

bool IKTwoBoneJob::Validate() const
{
    bool valid = true;
    valid &= start_joint && mid_joint && end_joint;
    valid &= start_joint_correction && mid_joint_correction;
    valid &= ozz::math::AreAllTrue1(ozz::math::IsNormalizedEst3(mid_axis));
    return valid;
}

namespace
{

// Local data structure used to share constant data accross ik stages.
struct IKConstantSetup {
    IKConstantSetup(const IKTwoBoneJob& _job)
    {
        // Prepares constants
        one       = simd_float4::one();
        mask_sign = simd_int4::mask_sign();
        m_one     = Xor(one, mask_sign);

        // Computes inverse matrices required to change to start and mid spaces.
        // If matrices aren't invertible, they'll be all 0 (ozz::math
        // implementation), which will result in identity correction quaternions.
        SimdInt4 invertible;
        (void)invertible;
        inv_start_joint              = Invert(*_job.start_joint, &invertible);
        const Float4x4 inv_mid_joint = Invert(*_job.mid_joint, &invertible);

        // Transform some positions to mid joint space (_ms)
        const SimdFloat4 start_ms =
            TransformPoint(inv_mid_joint, _job.start_joint->cols[3]);
        const SimdFloat4 end_ms =
            TransformPoint(inv_mid_joint, _job.end_joint->cols[3]);

        // Transform some positions to start joint space (_ss)
        const SimdFloat4 mid_ss =
            TransformPoint(inv_start_joint, _job.mid_joint->cols[3]);
        const SimdFloat4 end_ss =
            TransformPoint(inv_start_joint, _job.end_joint->cols[3]);

        // Computes bones vectors and length in mid and start spaces.
        // Start joint position will be treated as 0 because all joints are
        // expressed in start joint space.
        start_mid_ms                  = -start_ms;
        mid_end_ms                    = end_ms;
        start_mid_ss                  = mid_ss;
        const SimdFloat4 mid_end_ss   = end_ss - mid_ss;
        const SimdFloat4 start_end_ss = end_ss;
        start_mid_ss_len2             = Length3Sqr(start_mid_ss);
        mid_end_ss_len2               = Length3Sqr(mid_end_ss);
        start_end_ss_len2             = Length3Sqr(start_end_ss);
    }

    // Constants
    SimdFloat4 one;
    SimdFloat4 m_one;
    SimdInt4   mask_sign;

    // Inverse matrices
    Float4x4 inv_start_joint;

    // Bones vectors and length in mid and start spaces (_ms and _ss).
    SimdFloat4 start_mid_ms;
    SimdFloat4 mid_end_ms;
    SimdFloat4 start_mid_ss;
    SimdFloat4 start_mid_ss_len2;
    SimdFloat4 mid_end_ss_len2;
    SimdFloat4 start_end_ss_len2;
};

// Smoothen target position when it's further that a ratio of the joint chain
// length, and start to target length isn't 0.
// Inspired by http://www.softimageblog.com/archives/108
// and http://www.ryanjuckett.com/programming/analytic-two-bone-ik-in-2d/
bool SoftenTarget(const IKTwoBoneJob& _job, const IKConstantSetup& _setup, SimdFloat4* _start_target_ss, SimdFloat4* _start_target_ss_len2)
{
    // Hanlde position in start joint space (_ss)
    const SimdFloat4 start_target_original_ss =
        TransformPoint(_setup.inv_start_joint, _job.target);
    const SimdFloat4 start_target_original_ss_len2 =
        Length3Sqr(start_target_original_ss);
    const SimdFloat4 lengths =
        Sqrt(SetZ(SetY(_setup.start_mid_ss_len2, _setup.mid_end_ss_len2), start_target_original_ss_len2));
    const SimdFloat4 start_mid_ss_len             = lengths;
    const SimdFloat4 mid_end_ss_len               = SplatY(lengths);
    const SimdFloat4 start_target_original_ss_len = SplatZ(lengths);
    const SimdFloat4 bone_len_diff_abs =
        AndNot(start_mid_ss_len - mid_end_ss_len, _setup.mask_sign);
    const SimdFloat4 bones_chain_len = start_mid_ss_len + mid_end_ss_len;
    const SimdFloat4 da              = // da.yzw needs to be 0
        bones_chain_len *
        Clamp(simd_float4::zero(), simd_float4::LoadX(_job.soften), _setup.one);
    const SimdFloat4 ds = bones_chain_len - da;

    // Sotftens target position if it is further than a ratio (_soften) of the
    // whole bone chain length. Needs to check also that ds and
    // start_target_original_ss_len2 are != 0, because they're used as a
    // denominator.
    // x = start_target_original_ss_len > da
    // y = start_target_original_ss_len > 0
    // z = start_target_original_ss_len > bone_len_diff_abs
    // w = ds                           > 0
    const SimdFloat4 left      = SetW(start_target_original_ss_len, ds);
    const SimdFloat4 right     = SetZ(da, bone_len_diff_abs);
    const SimdInt4   comp      = CmpGt(left, right);
    const int        comp_mask = MoveMask(comp);

    // xyw all 1, z is untested.
    if ((comp_mask & 0xb) == 0xb)
    {
        // Finds interpolation ratio (aka alpha).
        const SimdFloat4 alpha = (start_target_original_ss_len - da) * RcpEstX(ds);
        // Approximate an exponential function with : 1-(3^4)/(alpha+3)^4
        // The derivative must be 1 for x = 0, and y must never exceeds 1.
        // Negative x aren't used.
        const SimdFloat4 three = simd_float4::Load1(3.f);
        const SimdFloat4 op    = SetY(three, alpha + three);
        const SimdFloat4 op2   = op * op;
        const SimdFloat4 op4   = op2 * op2;
        const SimdFloat4 ratio = op4 * RcpEstX(SplatY(op4));

        // Recomputes start_target_ss vector and length.
        const SimdFloat4 start_target_ss_len = da + ds - ds * ratio;
        *_start_target_ss_len2               = start_target_ss_len * start_target_ss_len;
        *_start_target_ss =
            start_target_original_ss *
            SplatX(start_target_ss_len * RcpEstX(start_target_original_ss_len));
    }
    else
    {
        *_start_target_ss      = start_target_original_ss;
        *_start_target_ss_len2 = start_target_original_ss_len2;
    }

    // The maximum distance we can reach is the soften bone chain length: da
    // (stored in !x). The minimum distance we can reach is the absolute value of
    // the difference of the 2 bone lengths, |d1-d2| (stored in z). x is 0 and z
    // is 1, yw are untested.
    return (comp_mask & 0x5) == 0x4;
}

SimdQuaternion ComputeMidJoint(const IKTwoBoneJob& _job, const IKConstantSetup& _setup, _SimdFloat4 _start_target_ss_len2)
{
    // Computes expected angle at mid_ss joint, using law of cosine (generalized
    // Pythagorean).
    // c^2 = a^2 + b^2 - 2ab cosC
    // cosC = (a^2 + b^2 - c^2) / 2ab
    // Computes both corrected and initial mid joint angles
    // cosine within a single SimdFloat4 (corrected is x component, initial is y).
    const SimdFloat4 start_mid_end_sum_ss_len2 =
        _setup.start_mid_ss_len2 + _setup.mid_end_ss_len2;
    const SimdFloat4 start_mid_end_ss_half_rlen =
        SplatX(simd_float4::Load1(.5f) * RSqrtEstXNR(_setup.start_mid_ss_len2 * _setup.mid_end_ss_len2));
    // Cos value needs to be clamped, as it will exit expected range if
    // start_target_ss_len2 is longer than the triangle can be (start_mid_ss +
    // mid_end_ss).
    const SimdFloat4 mid_cos_angles_unclamped =
        (SplatX(start_mid_end_sum_ss_len2) -
         SetY(_start_target_ss_len2, _setup.start_end_ss_len2)) *
        start_mid_end_ss_half_rlen;
    const SimdFloat4 mid_cos_angles =
        Clamp(_setup.m_one, mid_cos_angles_unclamped, _setup.one);

    // Computes corrected angle
    const SimdFloat4 mid_corrected_angle = ACosX(mid_cos_angles);

    // Computes initial angle.
    // The sign of this angle needs to be decided. It's considered negative if
    // mid-to-end joint is bent backward (mid_axis direction dictates valid
    // bent direction).
    const SimdFloat4 bent_side_ref  = Cross3(_setup.start_mid_ms, _job.mid_axis);
    const SimdInt4   bent_side_flip = SplatX(
        CmpLt(Dot3(bent_side_ref, _setup.mid_end_ms), simd_float4::zero())
    );
    const SimdFloat4 mid_initial_angle =
        Xor(ACosX(SplatY(mid_cos_angles)), And(bent_side_flip, _setup.mask_sign));

    // Finally deduces initial to corrected angle difference.
    const SimdFloat4 mid_angles_diff = mid_corrected_angle - mid_initial_angle;

    // Builds queternion.
    return SimdQuaternion::FromAxisAngle(_job.mid_axis, mid_angles_diff);
}

SimdQuaternion ComputeStartJoint(const IKTwoBoneJob& _job, const IKConstantSetup& _setup, const SimdQuaternion& _mid_rot_ms, _SimdFloat4 _start_target_ss, _SimdFloat4 _start_target_ss_len2)
{
    // Pole vector in start joint space (_ss)
    const SimdFloat4 pole_ss =
        TransformVector(_setup.inv_start_joint, _job.pole_vector);

    // start_mid_ss with quaternion mid_rot_ms applied.
    const SimdFloat4 mid_end_ss_final = TransformVector(
        _setup.inv_start_joint,
        TransformVector(*_job.mid_joint, TransformVector(_mid_rot_ms, _setup.mid_end_ms))
    );
    const SimdFloat4 start_end_ss_final = _setup.start_mid_ss + mid_end_ss_final;

    // Quaternion for rotating the effector onto the target
    const SimdQuaternion end_to_target_rot_ss =
        SimdQuaternion::FromVectors(start_end_ss_final, _start_target_ss);

    // Calculates rotate_plane_ss quaternion which aligns joint chain plane to
    // the reference plane (pole vector). This can only be computed if start
    // target axis is valid (not 0 length)
    // -------------------------------------------------
    SimdQuaternion start_rot_ss = end_to_target_rot_ss;
    if (AreAllTrue1(CmpGt(_start_target_ss_len2, simd_float4::zero())))
    {
        // Computes each plane normal.
        const ozz::math::SimdFloat4 ref_plane_normal_ss =
            Cross3(_start_target_ss, pole_ss);
        const ozz::math::SimdFloat4 ref_plane_normal_ss_len2 =
            ozz::math::Length3Sqr(ref_plane_normal_ss);
        // Computes joint chain plane normal, which is the same as mid joint axis
        // (same triangle).
        const ozz::math::SimdFloat4 mid_axis_ss =
            TransformVector(_setup.inv_start_joint, TransformVector(*_job.mid_joint, _job.mid_axis));
        const ozz::math::SimdFloat4 joint_plane_normal_ss =
            TransformVector(end_to_target_rot_ss, mid_axis_ss);
        const ozz::math::SimdFloat4 joint_plane_normal_ss_len2 =
            ozz::math::Length3Sqr(joint_plane_normal_ss);
        // Computes all reciprocal square roots at once.
        const SimdFloat4 rsqrts =
            RSqrtEstNR(SetZ(SetY(_start_target_ss_len2, ref_plane_normal_ss_len2), joint_plane_normal_ss_len2));

        // Computes angle cosine between the 2 normalized normals.
        const SimdFloat4 rotate_plane_cos_angle =
            ozz::math::Dot3(ref_plane_normal_ss * SplatY(rsqrts), joint_plane_normal_ss * SplatZ(rsqrts));

        // Computes rotation axis, which is either start_target_ss or
        // -start_target_ss depending on rotation direction.
        const SimdFloat4 rotate_plane_axis_ss = _start_target_ss * SplatX(rsqrts);
        const SimdFloat4 start_axis_flip =
            And(SplatX(Dot3(joint_plane_normal_ss, pole_ss)), _setup.mask_sign);
        const SimdFloat4 rotate_plane_axis_flipped_ss =
            Xor(rotate_plane_axis_ss, start_axis_flip);

        // Builds quaternion along rotation axis.
        const SimdQuaternion rotate_plane_ss = SimdQuaternion::FromAxisCosAngle(
            rotate_plane_axis_flipped_ss,
            Clamp(_setup.m_one, rotate_plane_cos_angle, _setup.one)
        );

        if (_job.twist_angle != 0.f)
        {
            // If a twist angle is provided, rotation angle is rotated along
            // rotation plane axis.
            const SimdQuaternion twist_ss = SimdQuaternion::FromAxisAngle(
                rotate_plane_axis_ss, simd_float4::Load1(_job.twist_angle)
            );
            start_rot_ss = twist_ss * rotate_plane_ss * end_to_target_rot_ss;
        }
        else
        {
            start_rot_ss = rotate_plane_ss * end_to_target_rot_ss;
        }
    }
    return start_rot_ss;
}

void WeightOutput(const IKTwoBoneJob& _job, const IKConstantSetup& _setup, const SimdQuaternion& _start_rot, const SimdQuaternion& _mid_rot)
{
    const SimdFloat4 zero = simd_float4::zero();

    // Fix up quaternions so w is always positive, which is required for NLerp
    // (with identity quaternion) to lerp the shortest path.
    const SimdFloat4 start_rot_fu =
        Xor(_start_rot.xyzw,
            And(_setup.mask_sign, CmpLt(SplatW(_start_rot.xyzw), zero)));
    const SimdFloat4 mid_rot_fu = Xor(
        _mid_rot.xyzw, And(_setup.mask_sign, CmpLt(SplatW(_mid_rot.xyzw), zero))
    );

    if (_job.weight < 1.f)
    {
        // NLerp start and mid joint rotations.
        const SimdFloat4 identity    = simd_float4::w_axis();
        const SimdFloat4 simd_weight = Max(zero, simd_float4::Load1(_job.weight));

        // Lerp
        const SimdFloat4 start_lerp = Lerp(identity, start_rot_fu, simd_weight);
        const SimdFloat4 mid_lerp   = Lerp(identity, mid_rot_fu, simd_weight);

        // Normalize
        const SimdFloat4 rsqrts =
            RSqrtEstNR(SetY(Length4Sqr(start_lerp), Length4Sqr(mid_lerp)));
        _job.start_joint_correction->xyzw = start_lerp * SplatX(rsqrts);
        _job.mid_joint_correction->xyzw   = mid_lerp * SplatY(rsqrts);
    }
    else
    {
        // Quatenions don't need interpolation
        _job.start_joint_correction->xyzw = start_rot_fu;
        _job.mid_joint_correction->xyzw   = mid_rot_fu;
    }
}
} // namespace

bool IKTwoBoneJob::Run() const
{
    if (!Validate())
    {
        return false;
    }

    // Early out if weight is 0.
    if (weight <= 0.f)
    {
        // No correction.
        *start_joint_correction = *mid_joint_correction =
            SimdQuaternion::identity();
        // Target isn't reached.
        if (reached)
        {
            *reached = false;
        }
        return true;
    }

    // Prepares constant ik data.
    const IKConstantSetup setup(*this);

    // Finds soften target position.
    SimdFloat4 start_target_ss;
    SimdFloat4 start_target_ss_len2;
    const bool lreached =
        SoftenTarget(*this, setup, &start_target_ss, &start_target_ss_len2);
    if (reached)
    {
        *reached = lreached && weight >= 1.f;
    }

    // Calculate mid_rot_local quaternion which solves for the mid_ss joint
    // rotation.
    const SimdQuaternion mid_rot_ms =
        ComputeMidJoint(*this, setup, start_target_ss_len2);

    // Calculates end_to_target_rot_ss quaternion which solves for effector
    // rotating onto the target.
    const SimdQuaternion start_rot_ss = ComputeStartJoint(
        *this, setup, mid_rot_ms, start_target_ss, start_target_ss_len2
    );

    // Finally apply weight and output quaternions.
    WeightOutput(*this, setup, start_rot_ss, mid_rot_ms);

    return true;
}
} // namespace animation
} // namespace ozz
